﻿#include <iostream>
#include <windows.h>
#include <shellapi.h>
#include <fstream>
#include <filesystem>
#include "encode.h"
#include "res/embedRes.hpp"
// By Zero123

#define PAUSE_RETURN(state) system("pause"); return state

// 检查是不是管理员模式运行
static bool checkRunAsAdmin()
{
    BOOL fIsRunAsAdmin = FALSE;
    PSID pAdminGroup = NULL;
    SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;

    // 建立一个管理员组 SID
    if (AllocateAndInitializeSid(&NtAuthority, 2,
        SECURITY_BUILTIN_DOMAIN_RID,
        DOMAIN_ALIAS_RID_ADMINS,
        0, 0, 0, 0, 0, 0,
        &pAdminGroup))
    {
        // 检查当前 token 是否在管理员组
        if (!CheckTokenMembership(NULL, pAdminGroup, &fIsRunAsAdmin))
        {
            fIsRunAsAdmin = FALSE;
        }
        FreeSid(pAdminGroup);
    }
    return fIsRunAsAdmin == TRUE;
}

// 以管理员模式运行自己
static bool runAsAdmin()
{
    char szPath[MAX_PATH];
    if (GetModuleFileNameA(NULL, szPath, MAX_PATH))
    {
        SHELLEXECUTEINFOA sei = { sizeof(sei) };
        sei.lpVerb = "runas";         // 请求以管理员权限运行
        sei.lpFile = szPath;          // 可执行文件路径
        sei.hwnd = NULL;
        sei.nShow = SW_NORMAL;

        if (!ShellExecuteExA(&sei))
        {
            DWORD dwError = GetLastError();
            if (dwError == ERROR_CANCELLED)
            {
                std::cerr << "用户拒绝了管理员权限请求。\n";
            }
            else
            {
                std::cerr << "提权失败，错误码: " << dwError << "\n";
            }
            return false;
        }
    }
    return true;
}

int main()
{
    if (!checkRunAsAdmin())
    {
        runAsAdmin();
        return 0;
    }
    Encoding::initEnvcode();
	std::cout << "自动化修复VC++运行库缺失问题 By Zero123\n\n";
    char systemDir[MAX_PATH];
    UINT len = GetSystemDirectoryA(systemDir, MAX_PATH);
    if (len == 0 || len > MAX_PATH)
    {
        std::cerr << "System32路径获取失败。\n";
        PAUSE_RETURN(1);
    }
	std::cout << "System32: " << systemDir << "\n";
	int missCount = 0;
    std::filesystem::path system32 { systemDir };
    for (const auto& [fileName, data] : EmbedRes::resourceMap)
    {
        auto dllPath = system32 / fileName;
        if(std::filesystem::is_regular_file(dllPath))
        {
            continue;
		}
        missCount++;
		std::cout << "缺失: " << fileName << "\n";
        std::ofstream outFile{ dllPath, std::ios::binary };
        if (!outFile)
        {
            std::cerr << "File Create Error: " << dllPath.string() << "\n";
            continue;
        }
        outFile.write(reinterpret_cast<const char*>(data.first), data.second);
        if (!outFile)
        {
            std::cerr << "File Write Error: " << dllPath.string() << "\n";
            continue;
        }
        outFile.close();
		std::cout << ">> 已补全: " << fileName << "\n";
	}
    if(missCount == 0)
    {
        std::cout << "未发现缺失的VC++运行库。\n";
	}
    else
    {
		std::cout << "共补全 " << missCount << " 个VC++运行库。\n";
    }
    PAUSE_RETURN(0);
}